python 🤝 `defer`

Bringing Go's defer statement to Python!

I'm a big fan of per-scope finalizers, à la Golangs's defer. It's super convenient for "fire-and-forget" cleanup, especially when dealing with expensive resources. Yes, python has try/except/finally blocks but I'm always forgetting the finally.. and yes I could use a context manager but all that extra indentation can get so clunky! Luckily, I've been playing around a bunch with @alexmojaki's excellent sorcery library, and the underlying executing, and it happens to be perfect for this particular brand of abuse.

Take 1: Explicit Function Calls

My first take was to write an explicit defer function. This would work similar to other attempts at bringing defer to python, but without the cumbersome decorators. Something like this:

def foo():
    instance = allocate_gpu()
    defer(instance.shutdown)
    
    results = train_a_big_model(instance)
    defer(lambda: print(f"Logging results no matter what: {results}")
    really_hope_this_doesnt_crash(instance)

instance.shutdown() should be called no matter what happens when this function runs.

By analyzing how perturbations of the source impact the AST, executing can identify exactly what node is being executed at a given frame. By combining this with sys.settrace(), we can catch the first call to defer, and create a function-local stack of finalizers. We then take that stack and wrap it in a child trace function, which we attach to the function via foo.f_trace.

Subsequent calls to defer then simply push finalizers to the stack. When we get the trace for exiting that function (via a return or a raise), we work our way down the stack, running each finalizer before allowing the function to conclude.

Turns out this works pretty well! You can see it implemented here. However, that ugly function call was really bugging me. I'd really love to not have to wrap all my statements in lambda: ... to use this.

Take 2: Live AST Rewriting & Compilation

What a great excuse to implement the world's tiniest JIT compiler. Ideally, we can just do something like this:

def foo():
    instance = allocate_gpu()
    instance.shutdown() in defer
    
    lets_go_gpt_5_time(instance)

where the in defer suffix is the most elegant way I could find to keep the parser happy.

This is cool; we no longer have to wrap our finalizer code in a lambda. We can just call it like normal, but be assured that line won't actually run until right before the function exits. Doesn't seem possible, right? Isn't this just going to call instance.shutdown() before I've had the chance to do anything?

Welcome to the magic of python metaprogramming. My initial idea was to use a tracer to intercept lines before they execute, replacing the instructions in-memory just before they're hit. Instead of calling instance.shutdown(), the bytecode line two would do something like defer(lambda: instance.shutdown()). Unfortunately—save for resorting to manually assembling the bytecode and doing some sketchy direct memory accesses—swapping the instructions out of a running python function doesn't seem possible.

What is possible is to replace the scope's reference to a function right after it's defined. Tracing each line, we can step in right after a FunctionDef node finishes executing. Using the name of the node, we can look up a reference to the compiled FunctionType in the current module. All that's left is to transform the AST as we described above, compile the new function, and swap out the __code__ property of the original.

The final logic looks something like this:

ast = RewriteDefers.transform(node)
compiled = compile(ast, frame.f_code.co_filename, "exec")

locals = frame.f_locals.copy()
exec(compiled, frame.f_globals, locals)

frame.f_locals[node.name].__code__ = locals[node.name].__code__

We now have a working, decorator-free defer in python.

Try it out!

You can use defer in either of the ways detailed above with a simple pip install python-defer . Check out the repo for more on getting started, and hit me up with any creative uses you come up with.